package org.example.cases.sf1;

import com.antgroup.geaflow.api.graph.compute.VertexCentricCompute;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricCombineFunction;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricComputeFunction;
import com.antgroup.geaflow.model.graph.edge.IEdge;
import com.antgroup.geaflow.model.graph.vertex.IVertex;
import org.example.function.AbstractVcFunc;

import java.util.Iterator;
import java.util.List;

public class Case3Algorithm {

    public static class Case3VCAlgorithm extends VertexCentricCompute<Long, Double, Double, Double> {

        public Case3VCAlgorithm(long iterations) {
            super(iterations);
        }

        @Override
        public VertexCentricComputeFunction<Long, Double, Double, Double> getComputeFunction() {
            return new Case3VertexCentricComputeFunction();
        }

        @Override
        public VertexCentricCombineFunction<Double> getCombineFunction() {
            return null;
        }
    }

    public static class Case3VertexCentricComputeFunction extends AbstractVcFunc<Long, Double, Double, Double> {

        @Override
        public void compute(Long vertexId, Iterator<Double> messageIterator) {
            IVertex<Long, Double> vertex = this.context.vertex().get();

            if (this.context.getIterationId() == 1) {
                List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
                if (!outEdges.isEmpty()) {
                    double outSum = 0.0;
                    for (IEdge<Long, Double> outEdge : outEdges) {
                        outSum += outEdge.getValue();
                        this.context.sendMessage(outEdge.getTargetId(), outEdge.getValue());
                    }
                    this.context.setNewVertexValue(-outSum);
                }
            } else {
                double inSum = 0.0;
                while (messageIterator.hasNext()) {
                    double value = messageIterator.next();
                    inSum += value;
                }
                Double outSum = vertex.getValue();
                if (outSum < 0) {
                    this.context.setNewVertexValue(-inSum / outSum);
                }
            }
        }
    }
}
